import math
import numpy
import pylab
import operator

class Connector:
    '''
    class to generate a connectivity matrix - you initialize the object
    with the  connectivity probs and the population sizes, 
    and then it can generate realizations of those probabilities for you, 
    in matrix form.
    '''

    def __init__(self, numE, numI, e2e, e2i, i2e, i2i, e2er, e2ir, i2er, i2ir):
	'''Creates a new instance with specified parameters
	@param numE: number of excitatory neurons (must be an even number)
	@param numI: number of inhibitory neurons (must be an even number)
	@param e2e: percent of excitatory to excitatory connections
	@param e2i: percent  of excitatory to inhibitory synapses 
	@param i2e: percent of inhib to excit synapse
	@param i2i: percent of inhib to inhib synapse
	@param e2er: percent of reciprocal synapse given a one-way synapse
	@param e2ir: percent of recip given one-way 
	@param i2er: percent of recip given one-way
	@param i2ir: percent of recip given 1-way
	
	'''
	if(numE % 2 != 0):
	    raise 'everything is easier if numE is divisible by two.'
	if(numI %2 != 0):
	    raise 'everything is easier if numI is divisible by two.'
	self.numE = numE
	self.numI = numI
	self.e2e = e2e / 100.0
	self.e2i = e2i / 100.0
	self.i2e = i2e / 100.0
	self.i2i = i2i / 100.0
	
	self.e2er = e2er / 100.0
	self.e2ir = e2ir / 100.0
	self.i2er = i2er / 100.0
	self.i2ir = i2ir / 100.0

    def makeMatrix(self, randSeed = None):
	'''
	we're going to make a matrix (mat) where mat[i,j] tells you the connection 
	between i and j.   for now it's binary but later i think it'll be the 
	'strength' of that connection, and eventually maybe even the max conductance.

	the first self.numE entries in the matrix are the excitatory cells, and the 
	remainder are the inhibitory cells.

	@param randSeed: the random generator is seeded each time this function is run. 
	if you supply a seed it will be used, otherwise None will be used and the generator
	will probably wind up seeded by the current clock time 
	(check your version of numpy.random.seed for documentation).
	'''

	#seed it:
	numpy.random.seed(randSeed)
	
	#first thing is to make a square matrix. excitatory cells first, then inhibitory cells.
	numCells = self.numE+self.numI
	mat = numpy.zeros(numCells**2, dtype=int).reshape(numCells, numCells)

	#first make the reciprocal connections between the same populations (e2e, i2i)
	#i'm deciding on the numbers here, rather than in the constructor, in case we decide
	#that we want per-trial or per-neuron variation, 
	#then these will become means that we sample around rather than hard numbers.
	#decide on an integer number first for how many of the other excites we synapse onto:
	numE2e = int(self.e2e * self.numE)
	#now make a round number of those reciprocal:
	numE2er = int(self.e2er * numE2e)
	#likewise for inhibitory:
	numI2i = int(self.i2i * self.numI)
	numI2ir = int(self.i2ir * numI2i)
	self.recipSamePop(mat, numE2er, 0, self.numE)
	self.recipSamePop(mat, numI2ir, self.numE, self.numE+self.numI)
	
	#print('numE2e: %d, numE2eR: %d, numI2i: %d, numI2ir: %d' %(numE2e, numE2er, numI2i, numI2ir))

	#then make the reciprocal connections across the two populations
	#do the percents to numbers thing again:
	numE2i = int(self.e2i * self.numE)
	numE2ir = int(self.e2ir * numE2i)
	numI2e = int(self.i2e * self.numI)
	numI2er = int(self.i2er * numI2e)
	#sanity check:
	#print('numE2i: %d, numE2ir: %d, numI2e: %d, numI2er: %d' % (numE2i, numE2ir, numI2e, numI2er))
	if(numI2er != numE2ir):
	    print('Warning: reciprocal connections between E and I dont match:')
	    print('numE2ir: %d, numI2er: %d --  using %d' % (numE2ir, numI2er, numE2ir))
	
	totalConnects = numE2ir*numE + numI2er*numI
	#recipCrossPop(self, mat, totalConnects, pop1Connects, pop2Connects pop1Bound, pop2Bound)
	self.recipCrossPop(  mat, totalConnects, numE2ir, numI2er, self.numE, self.numE+self.numI)
	#then fill in the rest. 
	self.oneWayConnect(mat, numE2e, 0, self.numE, 0, self.numE)
	self.oneWayConnect(mat, numE2i, 0, self.numE, self.numE, self.numI+self.numE)
	self.oneWayConnect(mat, numI2e, self.numE, self.numI+self.numE, 0, self.numE)
	self.oneWayConnect(mat, numI2i, self.numE, self.numI+self.numE, self.numE, self.numI+self.numE)

	return mat

	
    #------------------------------------------------------------------------------------------oneWayConnect
    def oneWayConnect(self, mat, numConnects, sourcePopLowInd, sourcePopHighInd, targetPopLowInd, targetPopHighInd):
	'''
	make one way connections, assuming that the reciprocal connections are already in place, and that
	none of these should create new reciprocal connections.
	@param mat: the connectivity matrix
	@param numConnects: the number of synapses from sourcePopulation onto targetPopulation
	@param sourcePopLowInd: start index of the source pop (inclusive)
	@param sourcePopHighInd: end index of the source pop (not inclusive)
	@param targetPopLowInd: start index of the target pop 
	@param targetPopHighInd: you can guess
	'''
	
	slow = sourcePopLowInd
	shi = sourcePopHighInd
	tlow = targetPopLowInd
	thi  = targetPopHighInd
	
	for i in xrange(slow, shi):
	    numSyn = numConnects
	    #first of all, there might already be synapses from us to someone else - we're 
	    #reciprocally connected.  decrement the number remaining to do accordingly
	    numAlready = sum(mat[i,j] for j in range(tlow,thi))
	    myRow = mat[i,tlow:thi]
	   
# print('myRow:'+repr(myRow))
	    numSyn -= numAlready
	    
	    #next thing is to not synapse onto anyone who already synapses onto us: 2 cases:
	    #case 1: it's reciprocal : we already synapse onto them, 
	    #so we already subtracted for that and shouldn't consider it further
	    #case 2: not reciprocal:  we already determined it's not reciprocal; we shouldn't change that

	    #get an array where the jth spot is 1 iff neuron j already synapses onto us
	    already = mat[tlow:thi,i]
	    #print('already:')
	    #print(already)
	    #what we want is to NOT synapse onto those guys, so we're going to logically invert:
	    notAlready = numpy.array(map(operator.__not__,already))
	    #make sure we don't synapse onto ourselves:
	    if(slow == tlow): #if it's the same pop
		#print('samePop')
		notAlready[i-slow] = False #set our own index to false, so we're not selected.
	    #now make a list of all possible indices of postsynaptic neurons:
	    potentials = numpy.array(range(tlow,thi))
	    #now because we have an array of bools, we can sub-select from the potentials :
	    potentials = potentials[notAlready]
	    #now this is the list of neurons we can synapse onto without breaking any constraints

	    #if this happens we'll get an error below when j goes out of bounds of perm, but this can't hurt.
	    if(len(potentials) < numSyn):
		print('its time to deal with the case of not having enough neurons to synapse onto')
		print('neuron %d, trying to get %d synapses out of this list:' % (i, numSyn))
		print(potentials)
		print('matrix:')
		print(mat)
	    
	    #get a random permutation of the indices of all the legal neurons
	    perm = numpy.random.permutation(potentials)
	    #print('numSyn: %s'%numSyn)
	    #print(perm)
	    #set the synapses from i to j, bringing i back to matrix coords
	    for j in xrange(numSyn):
		mat[i,perm[j]] = 1   #the j coord is already mapped into tlow,thi
		



    #--------------------------------------------------------------------------------------recipSamePop

    def recipSamePop(self, matrix, numConnects, lowInd, highInd):
	'''
	does the book-keeping to make a specified number of reciprocal connections 
	within a given population, eg excitatory-excitatory reciprocal synapses. 
	@parameter matrix: (reference to) the connectivity matrix we're building
	@parmaeter numConnects: the number of reciprocal connections each neuron should have
	@parameter lowInd: the start index of this population
	@parameter highInd: the end index of this pop
	'''

	#the size of the population (required to be an even number at the constructor:)
	num = highInd - lowInd
	#we'll keep a count of how many connections each neuron already has
	counts = numpy.zeros(num)
	halfNum = num/2
	target = numConnects*halfNum
	history = [] #everything before this happened before history started.
	lastPair = [] #we'll use this to make sure if we backtrack that we don't go down the same path.
        branchCounts = numpy.zeros(target) #we'll count how many times we return to a branch point to help us figure out when to backtrack.


	print('making a total of %d Intra-pop connections for popSize %d' % (numConnects, (highInd-lowInd)))

	#we're making numConnects*halfNum (=target) reciprocally connected pairs.
	#the problem is,  we might get most of the way done,  and then get 'stuck' in a situation
	#with no valid solution from that point,   eg if you make all but one pair, 
	#but in doing so you've filled the counts on every neuron except one neuron has count
	#equal to numConnects - 2,  but then you can't have an "autapse" SO
	#we need to be able to determine when we've hit a dead-end,  and backtrack from there. 
	#in order to backtrack,  we'll mantain a history of all the pairs we've connected, so that
	#we can undo them if they're not working. 
	while(len(history) < target):
            branch = len(history)
            branchCounts[branch] += 1
            #check to see if we've been spinning our wheels by returning to this branchpoint repeatedly:
            #halfNum - target is the number of pairs left that we need to make,
            #since this is non-deterministic, let's assume that 3 tries per pair is enough
            if(branchCounts[branch] > ((target+1) - branch)*2):
                #print('branch %d out of target: %d, count: %d' % (branch, target, branchCounts[branch]))
                branchCounts[branch] = 0
                lastPair = history.pop()
                self.backTrack(matrix, lastPair, lowInd, lowInd, counts, counts, 0)
                continue

	    #choose the first index at random, but keep looking if the one we choose is full:
	    #choosing randomly and then going incrementally is still 'random' and this guarantees
	    #that we find what we need in linear time if the constraints are actually hard to fit.
	    #this algorithm assumes that the constraints are 'sane', eg that the numbers are such that
	    #we can always choose a first neuron, and the hard part is in finding a valid partner.
	    i = numpy.random.randint(0, num)
	    while(counts[i] == numConnects):
		i = (i + 1) % num

	    #now do the same for the second index, also make sure i!=j
	    #also make sure there isn't already a synapse between them:
	    j = numpy.random.randint(0, num)
	    jCount = 0

	    #we are looking for a legal pair, which means j!=i, j isn't max'd on synapses, j and i aren't
	    #already connected,  j and i aren't a pair we're backtracking away from.     
	    #further, if there are no valid choices we don't want to loop eternally,  so we keep jCount
	    #and quit if that gets too high.
	    while(jCount < num and (j==i or counts[j] == numConnects or matrix[i+lowInd,j+lowInd] > 0 or 
				    (lastPair.count(i) and lastPair.count(j) ))  ):
		j = (j + 1) % num
		jCount += 1
		#print('j: %d', j)
		#print(j==i)
		#print(counts[j] == numConnects)
		#print(matrix[i+lowInd,j+lowInd] > 0)

	    #once we're out, check to see if we exited due to success,  or if the trouble is that we've 
	    #had a problem and need to backtrack:
	    if(jCount < num):
		#map back to the actual matrix coords instead of our local 0-rooted coords:
		matrix[i+lowInd,j+lowInd] = 1
		matrix[j+lowInd,i+lowInd] = 1
		#increment the counts.
		counts[i] += 1
		counts[j] += 1
		history.append([i,j]) #record in the annals of history
		lastPair = [] #this is empty unless we're in the process of backtracking
	    else:
		#backtrack:
                branchCounts[branch] = 0 #reset the branch counts if we're backtracking away
		lastPair = history.pop()
                self.backTrack(matrix, lastPair, lowInd, lowInd, counts, counts, 0)
		
		
    #---------------------------------------------------------------------------------------backTrack
    def backTrack(self, mat, ijPair, iMatOffset, jMatOffset,  iCounts, jCounts, jBaseOffset):
        '''
        common code for backtracking : decrement iCounts[i] and jCounts[j], 
        and set the relevant matrix entries using the offsets.
        '''
        i = ijPair[0]
        j = ijPair[1] - jBaseOffset
        iCounts[i] -= 1
        jCounts[j] -= 1
        mat[i+iMatOffset,j+jMatOffset] = 0
        mat[j+jMatOffset,i+iMatOffset] = 0
		

  
    #-------------------------------------------------------------------------------------recipCrossPop
    def recipCrossPop(self, mat, totalConnects, pop1Connects, pop2Connects, pop1Bound, pop2Bound):
	'''
	makes reciprocal connections across two populations
	@param mat: the matrix we're building
	@param numConnects: the number of synapses we're adding to each neuron
	@param pop1Bound: population 1 goes from index zero up to this index
	@param pop2Bound: pop 2 goes from pop1Bound up to this index.
	'''
	#just like in the same-population case, we'll keep counts.
	pop1Size = pop1Bound
	pop1Counts = numpy.zeros(pop1Size)
	pop2Size = pop2Bound - pop1Bound
	pop2Counts = numpy.zeros(pop2Size)
	history = []
	lastPair = []
        target = totalConnects/2
        branchCounts = numpy.zeros(target) #we'll count how many times we return to a branch point to help us figure out when to backtrack.

	print('making a total of %d Cross-Pop connections: - pop1: %d , pop2: %d for %d pop1 neurons and %d pop2 units' % 
	      (totalConnects,pop1Connects, pop2Connects, pop1Size, pop2Size))

	maxws = 0
	while(len(history) < target):


            branch = len(history)
            branchCounts[branch] += 1
            #print('branch:%d, count: %d'% (branch,branchCounts[branch]))
            #check to see if we've been spinning our wheels by returning to this branchpoint repeatedly:
            #halfNum - target is the number of pairs left that we need to make,
            #since this is non-deterministic, let's assume that 3 tries per pair is enough
            if(branchCounts[branch] > ((target+1) - branch)):
                #print('branch %d out of target: %d, count: %d' % (branch, target, branchCounts[branch]))
                branchCounts[branch] = 0
                lastPair = history.pop()
                #f = [0,0]
                #f[0] = lastPair[0]
                #f[1] = lastPair[1]- pop1Bound
                self.backTrack(mat, lastPair, 0, pop1Bound, pop1Counts, pop2Counts, pop1Bound)
                continue
                


	    #choose randomly and then search linearly from there if our first guess is full:
	    i = numpy.random.randint(0,pop1Size)
	    #ws = sum(pop1Counts)
	    #if(ws > maxws):
	#	maxws = ws
	    #print('sum pop1Counts: %d  , pop1Size*pop1Connects:  %d,   max pop1Counts: %d' % (ws, pop1Size*pop1Connects, maxws))

	    while(pop1Counts[i] == pop1Connects):
		i = (i+1) % pop1Size
	    #likewise:
	    j = numpy.random.randint(pop1Bound,pop2Bound)
	    jCount = 0
	    while(jCount < pop2Size and (pop2Counts[j-pop1Bound] == pop2Connects or 
					 mat[i,j] > 0 or
					 (lastPair.count(i) and lastPair.count(j)) )):
		jCount +=1
		j += 1
		if(j == pop2Bound): #this is just modulus mapped away from zero.
		    j = pop1Bound

	    if(jCount < pop2Size):
		#put it in the matrix , increment, and move on.
		mat[i,j] = 1
		mat[j, i] = 1
		pop1Counts[i] += 1
		pop2Counts[j - pop1Bound] += 1
		history.append([i,j])
		lastPair = []
	    else:
		#backtrack
		lastPair = history.pop()
                branchCounts[branch] = 0
                #for backtracking we actually need to map j differently.
                #f = [0,0]
                #f[0] = lastPair[0]
                #f[1] = lastPair[1]- pop1Bound
                self.backTrack(mat, lastPair, 0, pop1Bound, pop1Counts, pop2Counts, pop1Bound)
		#i = lastPair[0]
		#j = lastPair[1]
		#pop1Counts[i] -= 1
		#pop2Counts[j-pop1Bound] -=1
		#mat[i,j] = 0
		#mat[j,i] = 0
	

#-=---------------------------------------------------------------------------------------main

if __name__ == "__main__":
    
    numE = 82
    numI = 82
    e2e = 25
    e2i = 20
    i2e = 20
    i2i = 20
    e2er = 20
    e2ir = 70
    i2er = 70
    i2ir = 30
    c = Connector(numE, numI,  e2e, e2i, i2e, i2i, e2er, e2ir, i2er, i2ir)
    for i in range(2,4):
        print(str(i))
        m = c.makeMatrix(35)
    print(m)
    eLow = 0
    eHi = numE
    iLow = numE
    iHi = numE+numI
    

    numE2i =  int( (e2i/100.0)  * numE)
    numE2ir = int( (e2ir/100.0) * numE2i)
    numI2e =  int( (i2e/100.0)  * numI)
    numI2er = int( (i2er/100.0) * numI2e)
    numE2e =  int( (e2e/100.0)  * numE)
    numE2er = int( (e2er/100.0) * numE2e)
    numI2i =  int( (i2i/100.0)  * numI)
    numI2ir = int( (i2ir/100.0) * numI2i)
    
    numCells = numE+numI

    e2es = [sum(m[i,j] for j in range(0,numE)) for i in range(0,numE)]
    e2is = [sum(m[i,j] for j in range(numE,numCells)) for i in range(0,numE)]
    i2is = [sum(m[i,j] for j in range(numE,numCells)) for i in range(numE, numCells)]
    i2es = [sum(m[i,j] for j in range(0,numE)) for i in range(numE, numCells)]
    print(e2es)
    print(e2is)
    print(i2is)
    print(i2es)
